import sys, uos, utarfile, uerrno, verify, uzlib, ubinascii

def update(filename, md5sum=b''):
	fwtarname = '/data/firmware.tar'
	runningname = '/data/running'
	uos.chdir('/data')
	# check md5
	if (len(md5sum) != 16):
		# 从文件中获取
		try:
			md5file = open(filename + '.md5', 'rb')
			md5sum = ubinascii.unhexlify(md5file.read(32))
		except:
			print('md5 file read failed.')
			return False
	verify.md5Init()
	try:
		fw = open(filename, 'rb')
		while True:
			data = fw.read(512)
			if (len(data) == 0):
				break
			verify.md5Update(data)
	except:
		md5sum_f = b''
  
	md5sum_f = verify.md5Final()
	if (md5sum != md5sum_f):
		print('md5sum check failed.')
		return False
     
	# zlib
	try:
		fw = open(filename, 'rb')
		fwtar = uzlib.DecompIO(fw, 25)
		fwtarfile = open(fwtarname, 'wb')
		while True:
			buf=fwtar.read(512)
			if (len(buf) == 0):
				break
			fwtarfile.write(buf)
		fw.close()
		fwtarfile.close()
	except:
		return False

	# check a or b
	try:
		running = open(runningname, 'r')
		mainfolder = running.read(1)
		if (mainfolder == 'a'):
			mainfolder = 'b'
		else:
			mainfolder = 'a'
		running.close()
	except:
		mainfolder = 'a'

	newfold = '/data/' + mainfolder
 
	try:
		rmtree(newfold)
	except:
		test = 0
	uos.mkdir(newfold)
	uos.chdir(newfold)
 
	# tar
	t = utarfile.TarFile(fwtarname)
	for i in t:
		if i.type == utarfile.DIRTYPE:
			makedirs(i.name, exist_ok=True)
		else:
			f = t.extractfile(i)
			dest = open(i.name, "wb")
			copyfileobj(f, dest)
			dest.close()
	# 备份旧的，写入首次运行
	try:
		uos.rename(runningname, runningname + '.bak')
	except:
		pass

	running = open(runningname, 'w+')
	running.write(mainfolder)
	running.close()
	open('/data/first', 'wb').close()
 
	return True

def copyfileobj(src, dest, length=512):
	if hasattr(src, "readinto"):
		buf = bytearray(length)
		while True:
			sz = src.readinto(buf)
			if not sz:
				break
			if sz == length:
				dest.write(buf)
			else:
				b = memoryview(buf)[:sz]
				dest.write(b)
	else:
		while True:
			buf = src.read(length)
			if not buf:
				break
			dest.write(buf)

def makedirs(name, exist_ok=False):
	s = ""
	comps = name.split("/")
	if comps[-1] == "":
		comps.pop()
	for i, c in enumerate(comps):
		s += c
		try:
			uos.mkdir(s)
		except OSError as e:
			if e.args[0] != uerrno.EEXIST:
				raise
			if i == len(comps) - 1:
				if exist_ok:
					return
				raise e
		finally:
			s += "/"

def walk(top, topdown=True):
	files = []
	dirs = []
	for dirent in uos.ilistdir(top):
		mode = dirent[1]
		fname = fsdecode(dirent[0])
		if S_ISDIR(mode):
			if fname != "." and fname != "..":
				dirs.append(fname)
		else:
			files.append(fname)
	if topdown:
		yield top, dirs, files
	for d in dirs:
		yield from walk(top + "/" + d, topdown)
	if not topdown:
		yield top, dirs, files

def fsdecode(s):
	if type(s) is str:
		return s
	return str(s, "utf-8")

def S_ISDIR(mode):
	"""Return True if mode is from a directory."""
	return (mode & 0xF000) == 0x4000

def rmtree(top):
	for path, dirs, files in walk(top, False):
		for f in files:
			uos.remove(path + "/" + f)
		uos.rmdir(path)
